import imageio
from moviepy import VideoFileClip
from PIL import Image, ImageDraw, ImageFilter
import numpy as np
import time
import os


def crop_rounded_corners(image, radius):
    """
    将图像的四个角裁剪为圆角，并启用抗锯齿
    :param image: 输入的图像
    :param radius: 圆角的半径
    :return: 裁剪后的图像
    """
    # 创建一个与图像大小相同的透明掩码
    mask = Image.new('L', image.size, 0)
    draw = ImageDraw.Draw(mask)

    # 绘制一个圆角矩形作为掩码
    draw.rounded_rectangle((0, 0, image.width, image.height), radius, fill=255)

    # 对掩码进行模糊处理，实现抗锯齿
    mask = mask.filter(ImageFilter.SMOOTH_MORE)

    # 将掩码应用到图像的透明通道
    result = image.copy()
    result.putalpha(mask)
    return result


def crop_top_rounded_corners(image, radius):
    """
    仅裁剪图像的上边两个角为圆角（左上和右上），下边保持直角
    :param image:PIL.Image对象
    :param radius:圆角的半径（像素）
    :return:裁剪后的图像（带透明通道）
    """
    width, height = image.size
    mask = Image.new('L', (width, height), 0)  # 全黑的掩码
    draw = ImageDraw.Draw(mask)

    # 1.绘制一个矩形（覆盖整个图像）
    draw.rectangle((0, 0, width, height), fill=255)

    # 2.在左上角和右上角绘制黑色圆形（覆盖掉矩形的角）
    # 左上角（用黑色覆盖）
    draw.ellipse((0, 0, radius * 2, radius * 2), fill=0)
    # 右上角（用黑色覆盖）
    draw.ellipse((width - radius * 2, 0, width, radius * 2), fill=0)

    # 3.重新绘制一个只有上面两个角是圆形的形状
    # 绘制上半部分的圆角矩形（仅覆盖顶部）
    draw.rounded_rectangle((0, 0, width, radius * 2), radius, fill=255)

    # 4.抗锯齿（可选）
    mask = mask.filter(ImageFilter.SMOOTH_MORE)

    # 5.应用掩码到图像的透明通道
    image.putalpha(mask)
    # image.save('test1.png')
    return image


def get_inner_region(frame):
    """
    获取机框模板的内部区域（假设内部区域是透明的）
    :param frame: 机框模板图像
    :return: 内部区域的坐标和大小
    """
    alpha = np.array(frame.getchannel('A'))
    frame_self_width = 0
    start_width_flag = False
    for pixel in alpha[frame.height // 2]:
        # 获取每个像素点的透明度，如果pixel大于200，则认为不在透明区域，这机框横向宽度+1
        if pixel > 200:
            start_width_flag = True
            frame_self_width += 1
        else:
            if start_width_flag:
                break
            else:
                frame_self_width += 1

    frame_self_height = 0
    start_height_flag = False
    for pixel in alpha:
        if pixel in alpha:
            if pixel[(frame.width // 2) // 2] > 200:
                start_height_flag = True
                frame_self_height += 1
            else:
                if start_height_flag:
                    break
                else:
                    frame_self_height += 1
    return frame_self_width, frame_self_height, frame.width - 2 * frame_self_width, frame.height - 2 * frame_self_height


def get_inner_region_v2(frame):
    """
    获取机框模板的内部区域（假设内部区域是透明的）
    :param frame: 机框模板图像
    :return: 内部区域的坐标和大小
    """
    alpha = np.array(frame.getchannel('A'))
    pixels1 = alpha[frame.height // 2]
    x_offset, screen_width = count_screen_offset_width_height(pixels1)
    pixels2 = [pixel[(frame.width // 2) // 2] for pixel in alpha]
    y_offset, screen_height = count_screen_offset_width_height(pixels2)
    return x_offset, y_offset, screen_width, screen_height


def get_watch_frame_inner_region(frame):
    """
    获取手表模板的内部区域（假设内部区域是透明的）
    :param frame: 机框模板图像
    :return: 内部区域的属性
    """
    # alpha = np.array(frame.getchannel('A'))
    # x, y = 0, 0
    # flag1 = True
    # min_x, max_x = 0, 0
    # while True:
    #     if alpha[x][y] != 0:
    #         flag1 = False
    #         x += 1
    #         max_x = x
    #     else:
    #         if flag1:
    #             min_x = x
    #             x += 1
    #         else:
    #             break
    # print(min_x, max_x)
    # print((max_x - min_x) // 2)
    x, y, width, height = 80, 289, 316, 316
    return x, y, width, height


def count_screen_offset_width_height(pixels: np):
    """
    计算x轴 y轴的offset， 计算截图的宽高
    :param pixels: 像素列表
    :return:
    """
    left, right = 0, 0
    max_right = len(pixels)
    x_offset_flag = False  # x轴偏移量计算标识
    screen_width_flag = False  # 截图宽度计算标识
    x_offset = 0
    screen_width = 0
    while right <= max_right:
        if not screen_width_flag:  # 判断screen_width_flag是否为False,进入计算边框的宽度
            if x_offset_flag:
                if pixels[right] <= 200:
                    x_offset = right
                    screen_width_flag = True
                else:
                    left = right
                    right += 1
            else:
                if pixels[right] <= 200:
                    if left == right:
                        right += 1
                        left = right
                else:
                    x_offset_flag = True
                    right += 1
        else:  # 判断screen_width_flag是否为Ture,进入计算边截图的宽度
            if pixels[right] > 200:
                screen_width = right - left
                break
            else:
                right += 1
    return x_offset, screen_width


def resize_to_fit(screenshot, target_width, target_height, device):
    """
    调整截图的大小以适应目标区域，同时保持长宽比，并启用抗锯齿
    :param screenshot: 截图
    :param target_width: 目标宽度
    :param target_height: 目标高度
    :param device: 设备类型
    """
    new_width = target_width
    new_height = target_height
    # 启用抗锯齿
    return screenshot.resize((new_width, new_height), Image.LANCZOS)


def add_frame_to_image(screenshot, frame, top_path, device, inner_region):
    """
    将单帧图像添加到机框
    :param screenshot:单帧截图
    :param frame: 机框模板
    :param top_path: 状态栏模板
    :param device: 设备类型
    :param inner_region: 机框内部区域参数（x_offset, y_offset, screen_width, screen_height）
    :return: 处理后的图像
    """
    inner_width = inner_region[2]
    inner_height = inner_region[3]

    # 调整截图大小以适应机框内部区域
    resized_screenshot = resize_to_fit(screenshot, inner_width, inner_height, device)

    # 裁剪截图的四个角为圆角
    resized_screenshot = crop_rounded_corners(resized_screenshot, radius=device.CORNER_RADIUS)

    # 创建一个机框模板大小相同的空白图像
    combined_image = Image.new('RGBA', frame.size)

    # 计算截图在内部区域的居中位置
    x_offset = inner_region[0]
    y_offset = inner_region[1]

    # 将截图粘贴到机框模板的内部区域
    combined_image.paste(resized_screenshot, (x_offset, y_offset), resized_screenshot)

    if top_path is not None:
        top = Image.open(top_path).convert('RGBA')
        # 顶部状态栏宽高比
        top_ratio = top.width / top.height
        top = top.resize((inner_width, int(inner_width / top_ratio)), Image.LANCZOS)
        combined_image.paste(top, (x_offset, y_offset), top)
        # 判断是否是沉浸式
        if 'immersive' in top_path:
            bottom_path = top_path.replace('top', 'bottom')
            bottom = Image.open(bottom_path).convert('RGBA')
            bottom_ratio = bottom.width / bottom.height
            bottom = bottom.resize((inner_width, int(inner_width / bottom_ratio)), Image.LANCZOS)
            combined_image.paste(bottom,
                                 (x_offset, y_offset + inner_region[3] - bottom.height + device.BOTTOM_Y_OFFSET),
                                 bottom)
    # 将机框模板粘贴到上层
    combined_image.paste(frame, (0, 0), frame)
    combined_image = combined_image.resize(device.DPI)

    return combined_image


def crop_to_circle_advanced(img, r):
    """将图片裁剪为圆形，非圆形区域透明"""
    width, height = img.size
    img = img.resize((2 * r, 2 * r), Image.LANCZOS)
    mask = Image.new("L", (2 * r, 2 * r), 0)
    draw = ImageDraw.Draw(mask)
    draw.ellipse((0, 0, 2 * r, 2 * r), fill=255)

    # 创建新图像，将非圆形区域设为透明
    result = Image.new("RGBA", (2 * r, 2 * r), (0, 0, 0, 0))
    result.paste(img, (0, 0), mask)
    return result


def add_watch_frame_to_image(screenshot, frame, x_offset, y_offset, r):
    """
    将单帧图像添加到机框
    :param screenshot:单帧截图
    :param frame: 机框模板
    :param x_offset: 图片x坐标
    :param y_offset: 图片y坐标
    :param r: 表盘圆半径
    :return: 处理后的图像
    """
    # 将截图裁剪为圆形
    resized_screenshot = crop_to_circle_advanced(screenshot, r)
    resized_screenshot.save('2.png')
    # 创建一个机框模板大小相同的空白图像
    combined_image = Image.new('RGBA', frame.size)
    # 将截图粘贴到机框模板的内部区域
    combined_image.paste(resized_screenshot, (x_offset, y_offset), resized_screenshot)
    # 将机框模板粘贴到上层
    combined_image.paste(frame, (0, 0), frame)
    return combined_image


def add_frame_to_video(video_path, frame_path, top_path, output_webp_path, device, progress_dialog=None):
    """
    给视频添加机框并到处为GIF
    :param video_path: 视频路径
    :param frame_path: 机框模板路径
    :param top_path: 状态栏模板路径
    :param output_webp_path: 输出webp路径
    :param device: 设备类型
    :param progress_dialog: 进度条控件
    :return:
    """
    # 加载机框模板
    frame = Image.open(frame_path).convert('RGBA')
    # 读取视频
    video = VideoFileClip(video_path)
    fps = video.fps
    # 处理每一帧
    processed_frames = []
    all_processed = int(video.fps * video.duration)
    current_frames_index = 1
    # 获取机框的inner_region
    frame_inner_region = get_inner_region_v2(frame)
    dpi_switch_flag = False  # dpi高宽切换标志
    for frame_image in video.iter_frames():
        # 将帧转换为PIL图像
        pil_frame = Image.fromarray(frame_image)
        if device.type == 'phone' and pil_frame.width > pil_frame.height:
            frame = frame.transpose(Image.ROTATE_90)  # 机框旋转90°
            if not dpi_switch_flag:
                device.DPI = (device.DPI[1], device.DPI[0])  # DPI高宽互换
                dpi_switch_flag = True
        # 添加机框
        processed_frame = add_frame_to_image(pil_frame, frame, top_path, device, frame_inner_region)

        # 将处理后的帧转换回numpy数组
        processed_frame = processed_frame.resize(device.DPI)
        processed_frames.append(np.array(processed_frame))

        # 处理进度条显示
        progress_dialog.setValue(int((current_frames_index / all_processed) * 99))
        current_frames_index += 1

        # 检查用户是否点击了取消按钮
        if progress_dialog.wasCanceled():
            break

    # 使用imageio到处webp
    imageio.mimsave(output_webp_path, processed_frames, format='WEBP', fps=fps, lossless=False, guantizer=90,
                    mode='RGBA', minimize=True)
    progress_dialog.setValue(100)


def wait_until_file_ready(filename, timeout=10):
    """等待文件完全写入"""
    start_time = time.time()
    last_size = -1

    while time.time() - start_time < timeout:
        current_size = os.path.getsize(filename)
        if current_size == last_size and current_size > 0:
            return True  # 文件大小稳定，可能已完成写入
        last_size = current_size
        time.sleep(0.1)

    return False


def add_watch_frame_to_video(video_path, frame_path, output_webp_path, progress_dialog=None):
    """
    给视频添加机框并到处为GIF
    :param video_path: 视频路径
    :param frame_path: 机框模板路径
    :param output_webp_path: 输出webp路径
    :param progress_dialog: 进度条控件
    :return:
    """
    # 加载机框模板
    frame = Image.open(frame_path).convert('RGBA')
    # 读取视频
    video = VideoFileClip(video_path)
    fps = video.fps
    # 处理每一帧
    processed_frames = []
    all_processed = int(video.fps * video.duration)
    current_frames_index = 1
    # 获取机框的inner_region
    x, y, width, height = get_watch_frame_inner_region(frame)
    for frame_image in video.iter_frames():
        # 将帧转换为PIL图像
        pil_frame = Image.fromarray(frame_image)

        # 添加机框
        processed_frame = add_watch_frame_to_image(pil_frame, frame, x, y, width)

        # 将处理后的帧转换回numpy数组
        processed_frames.append(np.array(processed_frame))

        # 处理进度条显示
        progress_dialog.setValue(int((current_frames_index / all_processed) * 99))
        current_frames_index += 1

        # 检查用户是否点击了取消按钮
        if progress_dialog.wasCanceled():
            break

    # 使用imageio到处webp
    imageio.mimsave(output_webp_path, processed_frames, format='WEBP', fps=fps, lossless=False, guantizer=90,
                    mode='RGBA', minimize=True)
    progress_dialog.setValue(100)


if __name__ == '__main__':
    watch_frame = 'resource/watch/机框.png'
    watch_frame = Image.open(watch_frame).convert('RGBA')
    get_watch_frame_inner_region(watch_frame)
